package com.ruoyi.websocket;




import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.JsonNode;

import com.ruoyi.aiConversation.domain.AiConversation;
import com.ruoyi.aiConversation.service.IAiConversationService;
import com.ruoyi.message.domain.Message;
import com.ruoyi.message.service.IMessageService;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.*;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
@Component
public class ChatWebSocketHandler implements WebSocketHandler {
    // Jackson ObjectMapper用于处理JSON
    private final ObjectMapper objectMapper = new ObjectMapper();
    Map<String, WebSocketConnect> connects = new ConcurrentHashMap<>();
    @Autowired
    private IAiConversationService aiConversationService;

    @Autowired
    private IMessageService messageService;

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        String sessionId = session.getId();
        if(session.getUri().getQuery().split("=").length > 2) {
            sessionId = session.getUri().getQuery().split("=")[2];
            System.out.println("继续使用旧的会话ID，Session ID: " + sessionId);
            // 创建新的WebSocketConnect对象
            WebSocketConnect connect = new WebSocketConnect(sessionId);
            connects.put(sessionId, connect);
        }
        else {
            Long userId = Long.valueOf(session.getUri().getQuery().split("=")[1]);
            System.out.println("新的WebSocket连接已建立，Session ID: " + session.getId());
            System.out.println(session.getUri());
            WebSocketConnect connect = new WebSocketConnect(sessionId);
            connects.put(sessionId, connect);
            AiConversation conversation = new AiConversation();
            conversation.setUserId(userId);
            conversation.setAiConversationId(sessionId);
            System.out.println(conversation);
            aiConversationService.insertAiConversation(conversation);
        }
    }
    @Override
    public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) {
        try {
            String sessionId;
            if(session.getUri().getQuery().split("=").length > 2) {
                sessionId = session.getUri().getQuery().split("=")[2];
            } else {
                sessionId = session.getId();
            }
            
            // 获取原始payload
            String payload = (String) message.getPayload();
            
            // 首先解析外层JSON，获取message字段
            JsonNode rootNode = objectMapper.readTree(payload);
            String innerMessage = rootNode.get("message").asText();
            
            // 解析内层JSON，获取实际的消息内容
            JsonNode innerNode = objectMapper.readTree(innerMessage);
            String actualMessage = innerNode.get("message").asText();
            
            // 保存用户消息
            Message message1 = new Message();
            message1.setMessageContent(actualMessage);  // 保存实际的消息内容
            message1.setAiConversationId(sessionId);
            message1.setSender("2");
            messageService.insertMessage(message1);
            
            // 处理回复
            WebSocketConnect connect = connects.get(sessionId);
            String reply = connect.reply(actualMessage);  // 使用实际的消息内容生成回复
            
            // 保存AI回复
            Message message2 = new Message();
            message2.setMessageContent(reply);
            message2.setSender("0");
            message2.setAiConversationId(sessionId);
            messageService.insertMessage(message2);
            
            // 构造响应消息
            WebMessage responseMessage = new WebMessage();
            responseMessage.setSession(sessionId);
            responseMessage.setMessage(reply);
            
            // 发送响应
            String responsePayload = objectMapper.writeValueAsString(responseMessage);
            session.sendMessage(new TextMessage(responsePayload));
            
        } catch (Exception e) {
            e.printStackTrace();
            try {
                session.sendMessage(new TextMessage("{\"error\": \"Invalid message format.\"}"));
            } catch (Exception ex) {
                ex.printStackTrace();
            }
        }
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        System.out.println("WebSocket传输错误：" + exception.getMessage());
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) throws Exception {
        String sessionId ;
        if(session.getUri().getQuery().split("=").length > 2) {
            sessionId = session.getUri().getQuery().split("=")[2];
        } else {
            sessionId = session.getId();
        }
        System.out.println("WebSocket连接关闭，Session ID: " + sessionId);
        // 从连接列表中移除已关闭连接的session ID
        connects.remove(sessionId);
        System.out.println("连接已移除，当前连接数: " + connects.size());
    }

    @Override
    public boolean supportsPartialMessages() {
        // 不支持消息分片
        return false;
    }
}
